import os
from langchain.agents import AgentType
from langchain.agents import load_tools, initialize_agent
from langchain.agents.tools import Tool
from langchain.chat_models import ChatOpenAI
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_community.chat_models import ChatTongyi, ChatOpenAI
from langchain.chains.conversation.memory import ConversationBufferMemory
import selftool
from vectorstore import load_vector_store
from langchain_community.embeddings import DashScopeEmbeddings

import config

os.environ["DASHSCOPE_API_KEY"] = "sk-53b4fbc8f87b4362b110e47e5bbf9b22"
os.environ["OPENAI_API_KEY"] = "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"


def CreateLLM(modelname: str):
    if modelname.count("qwen") > 0:
        return ChatTongyi(model_name=modelname)
    else:
        return ChatOpenAI(model_name=modelname)


class LLMChatBot():

    def __init__(self, module_name):
        self.llm = CreateLLM(module_name)

    def Chat(self, msg):
        try:
            return self.llm(msg).content
        except Exception as e:
            return "抱歉,尝试回答问题时发生了错误"


class AgentChatBot():
    def __init__(self, module_name: str):

        self.llm = CreateLLM(module_name)
        self.mytools = self.CreateMyToos()
        self.agent = self.CreateAgent()


    def CreateAgent(self):
        memory = ConversationBufferMemory(memory_key="chat_history")
        agent = initialize_agent(
            self.mytools,
            self.llm,
            agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
            verbose=True,
            memory=memory,
            handle_parsing_errors=True)
        return agent

    def CreateMyToos(self) -> list[Tool]:
        tools = [selftool.weather_tool, selftool.ProductInfo_tool, selftool.PartDefineInfo_tool]
        return tools

    def Chat(self, msg: str, ChatHistory: ChatMessageHistory):
        try:
            self.agent.memory.chat_memory = ChatHistory
            return self.agent(msg)["output"]
        except Exception as e:
            return "抱歉,尝试回答问题时发生了错误"


class KnowageBaseChatBot():


    def __init__(self, module_name: str, collection_name: str):
        self.llm = CreateLLM(module_name)
        self.collect_name = collection_name
        self.chain = self.createQaChain()

    def createQaChain(self):
        qdrant = load_vector_store(config.db_path, DashScopeEmbeddings(),
                                   collection_name=self.collect_name)
        qdrant.client.close()
        from langchain.chains import RetrievalQA
        qa_chain = RetrievalQA.from_chain_type(self.llm, retriever=qdrant.as_retriever())
        qdrant.client.close()
        return qa_chain
    def Chat(self, question, chat_history):
        return self.chain.run(question)
